# -*- coding: utf-8 -*-
# FileName:     model.py
# time:         23/4/6 006 下午 8:37
# Author:       Zhou Hang
# Description:  I don't want to write

import argparse
import sys
import time
import random
from pathlib import Path

import cv2
import numpy as np
import torch
from torch.backends import cudnn

from Yolo.models.experimental import attempt_load
from Yolo.utils.datasets import letterbox, LoadStreams, LoadImages
from Yolo.utils.general import check_img_size, non_max_suppression, scale_coords, increment_path, set_logging, \
    apply_classifier, check_imshow
from Yolo.utils.plots import plot_one_box2
from Yolo.utils.torch_utils import select_device, load_classifier, time_synchronized


class YoloModel:
    def __init__(self, weight_path='./Yolo/weight/yolov5s.pt',
                 conf_thres=0.25, iou_thres=0.45, device=''):
        self.init_model(weight_path, conf_thres, iou_thres, device)

    def init_model(self, weight_path, conf_thres, iou_thres, device):
        parser = argparse.ArgumentParser()
        parser.add_argument('--weights', nargs='+', type=str, default=weight_path, help='model.pt path(s)')
        parser.add_argument('--source', type=str, default='0', help='source')  # file/folder, 0 for webcam
        parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
        parser.add_argument('--conf-thres', type=float, default=conf_thres,
                            help='object confidence threshold')
        parser.add_argument('--iou-thres', type=float, default=iou_thres,
                            help='IOU threshold for NMS')
        parser.add_argument('--device', default=device,
                            help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
        parser.add_argument('--view-img', action='store_true', help='display results', default=False)
        parser.add_argument('--save-txt', action='store_true', help='save results to *.txt', default=False)
        parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels',
                            default=False)
        parser.add_argument('--nosave', action='store_true', help='do not save images/videos', default=True)
        parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
        parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS',
                            default=True)
        parser.add_argument('--augment', action='store_true', help='augmented inference',
                            default=True)
        parser.add_argument('--update', action='store_true', help='update all models')
        parser.add_argument('--project', default='runs/detect', help='save results to project/name')
        parser.add_argument('--name', default='exp', help='save results to project/name')
        parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
        self.opt = parser.parse_args()
        print(self.opt)
        source, weights, view_img, save_txt, imgsz = self.opt.source, self.opt.weights, self.opt.view_img, self.opt.save_txt, self.opt.img_size

        self.device = select_device(self.opt.device)
        self.half = self.device.type != 'cpu'  # half precision only supported on CUDA
        cudnn.benchmark = True

        # Load model
        self.model = attempt_load(weights, map_location=self.device)  # load FP32 model
        self.stride = int(self.model.stride.max())  # model stride
        self.imgsz = check_img_size(imgsz, s=self.stride)  # check img_size
        if self.half:
            self.model.half()  # to FP16

        # Get names and colors
        self.names = self.model.module.names if hasattr(self.model, 'module') else self.model.names
        self.colors = [[random.randint(0, 255) for _ in range(3)] for _ in self.names]

    def detect_img(self, name_list, img):
        """
        :param name_list: 文件名列表
        :param img: 待检测图片, 是一个cv2.imread或cv2.cap.read得到的图像信息
        :return: info_show: 检测输出的文字信息，处理后的图片直接通过img返回了
        """
        showimg = img
        with torch.no_grad():
            img = letterbox(img, new_shape=self.opt.img_size)[0]
            # Convert
            img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
            img = np.ascontiguousarray(img)
            img = torch.from_numpy(img).to(self.device)
            img = img.half() if self.half else img.float()  # uint8 to fp16/32
            img /= 255.0  # 0 - 255 to 0.0 - 1.0
            if img.ndimension() == 3:
                img = img.unsqueeze(0)
            # Inference
            pred = self.model(img, augment=self.opt.augment)[0]
            # Apply NMS
            pred = non_max_suppression(pred, self.opt.conf_thres,
                                       self.opt.iou_thres,
                                       classes=self.opt.classes,
                                       agnostic=self.opt.agnostic_nms)
            info_show = ""
            # Process detections
            for i, det in enumerate(pred):
                if det is not None and len(det):
                    # Rescale boxes from img_size to im0 size
                    det[:, :4] = scale_coords(img.shape[2:], det[:, :4], showimg.shape).round()
                    for *xyxy, conf, cls in reversed(det):
                        label = '%s %.2f' % (self.names[int(cls)], conf)
                        name_list.append(self.names[int(cls)])
                        single_info = plot_one_box2(xyxy,
                                                    showimg,
                                                    label=label,
                                                    color=self.colors[int(cls)],
                                                    line_thickness=2)
                        # print(single_info)
                        info_show = info_show + single_info + "\n"
        return info_show

    def detect(self, img_path: str, save_path: str):
        """
        通过img_path读取图片, 保存到save_path中，两个值可以相同
        """
        img = cv2.imread(img_path)
        if img is None:
            print('图片不存在')
            return
        self.detect_img([], img)
        cv2.imwrite(save_path, img)


def main():
    model = YoloModel()
    model.detect('./data/Image/10114.jpg', "./runs/test.jpg")


if __name__ == "__main__":
    main()
